clc;
clear;

tic;
startTime = cputime;
pde = process_pde_t_2d(3);
vector_h = [1/4 1/8 1/16];
% vector_dt = [1/8 1/32 1/64 1/256];
basis_type = "P2";
theta = 1;
disp(["basis_type: ", basis_type]);
if 0 == theta
    disp(["temporal_type: ", "forward Euler"]);
    error("stop");
elseif 1 == theta
    disp(["temporal_type: ", "backward Euler"]);
elseif 0.5 == theta
    disp(["temporal_type: ", "Crank-Nicolson"]);
end

u_Linf = zeros(size(vector_h,2),1);
u_L2 = zeros(size(vector_h,2),1);
u_H1 = zeros(size(vector_h,2),1);
p_Linf = zeros(size(vector_h,2),1);
p_L2 = zeros(size(vector_h,2),1);
p_H1 = zeros(size(vector_h,2),1);

for i = 1:size(vector_h,2)
    h = vector_h(i);
    mesh = process_mesh_h_2d(pde, h, basis_type);
    O1 = sparse(mesh.Np, mesh.Np);
    O2 = sparse(mesh.Nb, mesh.Np);
    O3 = sparse(mesh.Nb, mesh.Nb);
    O4 = zeros(mesh.Np,1);
    A1 = assemble_stiffness_matrix_2d(mesh, "P2", "dx", "P2", "dx", pde.nu);
    A2 = assemble_stiffness_matrix_2d(mesh, "P2", "dy", "P2", "dy", pde.nu);
    A3 = assemble_stiffness_matrix_2d(mesh, "P2", "dx", "P2", "dy", pde.nu);
    A5 = assemble_stiffness_matrix_2d(mesh, "P1", "x", "P2", "dx", -1);
    A6 = assemble_stiffness_matrix_2d(mesh, "P1", "x", "P2", "dy", -1);
    A = [2*A1+A2, A3, A5;
        A3', 2*A2+A1, A6;
        A5', A6', O1];
    Me = assemble_mass_matrix_2d(mesh);
    M = [Me, O3, O2;
        O3, Me, O2;
        O2', O2', O1];
    U1 = [pde.u1(mesh.node(:,1),mesh.node(:,2),0);
        pde.u2(mesh.node(:,1),mesh.node(:,2),0);
        pde.p(mesh.node(1:mesh.Np,1),mesh.node(1:mesh.Np,2),0)];
    if 1 == theta
        dt = 8*h^3;
    elseif 0.5 == theta
        % dt = vector_dt(i);
        dt = sqrt(h^3);
    end
    Nt = (pde.domain(6) - pde.domain(5))/dt;
    A(2*mesh.Nb+1,:) = 0;
    A(2*mesh.Nb+1,2*mesh.Nb+1) = 1;
    for m = 0:Nt-1
        t1 = m * dt;
        t2 = (m+1) * dt;
        [b1, b2] = assemble_load_vector_2d(mesh, pde, t1);
        bt1 = [b1; b2; O4];
        [b1, b2] = assemble_load_vector_2d(mesh, pde, t2);
        bt2 = [b1; b2; O4];
        [A, bt1] = process_boundary_t_2d(A, bt1, mesh, pde, t1);
        bt1(2*mesh.Nb+1,1) = pde.p(pde.domain(1), pde.domain(3), t1);
        % A(end,:) = 0;
        % A(end,end) = 1;
        % bt1(end,1) = pde.p(pde.domain(2), pde.domain(4), t1);
        [A, bt2] = process_boundary_t_2d(A, bt2, mesh, pde, t2);
        bt2(2*mesh.Nb+1,1) = pde.p(pde.domain(1), pde.domain(3), t2);
        % bt2(end,1) = pde.p(pde.domain(2), pde.domain(4), t1);

        % A_hat = M/(theta*dt) + A;
        % b_hat = theta*bt2 + (1-theta)*bt1 + (M/(theta*dt))*U1;
        % U_hat = A_hat\b_hat;
        % U2 = (U_hat - U1)/theta + U1;
        % U1 = U2;
        A_hat = M/dt + theta*A;
        b_hat = theta*bt2 + (1-theta)*bt1 + (M/dt - (1-theta)*A)*U1;
        U2 = A_hat\b_hat;
        U1 = U2;
    end

    u_Linf(i,1) = process_error_u_t_2d(mesh, pde, U2, "Linf");
    u_L2(i,1) = process_error_u_t_2d(mesh, pde, U2, "L2");
    u_H1(i,1) = process_error_u_t_2d(mesh, pde, U2, "H1");
    p_Linf(i,1) = process_error_p_t_2d(mesh, pde, U2, "Linf");
    p_L2(i,1) = process_error_p_t_2d(mesh, pde, U2, "L2");
    p_H1(i,1) = process_error_p_t_2d(mesh, pde, U2, "H1");

    fprintf("已完成\t%d/%d\n", i, size(vector_h,2));
end

process_convergence_h_2d(u_Linf, vector_h);
process_convergence_h_2d(u_L2, vector_h);
process_convergence_h_2d(u_H1, vector_h);
process_convergence_h_2d(p_Linf, vector_h);
process_convergence_h_2d(p_L2, vector_h);
process_convergence_h_2d(p_H1, vector_h);

elapsedTOCTime = toc;
elapsedCPUTime = cputime - startTime;
disp(["TOC time(s)", num2str(elapsedTOCTime)]);
disp(["CPU time(s)", num2str(elapsedCPUTime)]);


%% assemble_mass_matrix_2d
function matrix_M = assemble_mass_matrix_2d(mesh)
Nb1 = mesh.Nb;
Nb2 = Nb1;
Nlb1 = mesh.Nlb;
Nlb2 = Nlb1;
rows = zeros(2*Nlb1*Nlb2,1);
cols = rows;
values = rows;
flag = 1;
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), mesh.basis_type, "x");
    for i = 1:Nlb2
        for j = 1:Nlb1
            f = phi(:,j) .* phi(:,i);
            r = gauss_weights_loc * f;
            rows(flag) = mesh.elem(n,i);
            cols(flag) = mesh.elem(n,j);
            values(flag) = r;
            flag = flag + 1;
        end
    end
end
matrix_M = sparse(rows, cols, values, Nb2, Nb1);
end

%% assemble_stiffness_matrix_2d
function matrix_A = assemble_stiffness_matrix_2d(mesh, B1_type, D1_type, B2_type, D2_type, cor)
basis_map = containers.Map(...
    {'P1', 'P1b', 'P2'},...
    {struct('Nb', mesh.Np, 'Nlb',3),...
    struct('Nb', mesh.Nb, 'Nlb',4),...
    struct('Nb', mesh.Nb, 'Nlb',6)});
B1 = basis_map(B1_type);
Nb1 = B1.Nb();
Nlb1 = B1.Nlb();
B2 = basis_map(B2_type);
Nb2 = B2.Nb();
Nlb2 = B2.Nlb();
rows = zeros(2*Nlb1*Nlb2,1);
cols = rows;
values = rows;
flag = 1;
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi1 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B1_type, D1_type);
    phi2 = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), B2_type, D2_type);
    for i = 1:Nlb2
        for j = 1:Nlb1
            f = cor * phi1(:,j) .* phi2(:,i);
            r = gauss_weights_loc * f;
            rows(flag) = mesh.elem(n,i);
            cols(flag) = mesh.elem(n,j);
            values(flag) = r;
            flag = flag + 1;
        end
    end
end
matrix_A = sparse(rows, cols, values, Nb2, Nb1);
end

%% assemble_load_vector_2d
function [vector_b1, vector_b2] = assemble_load_vector_2d(mesh, pde, t)
vector_b1 = zeros(mesh.Nb,1);
vector_b2 = vector_b1;
for n = 1:mesh.Ne
    E = mesh.node(mesh.elem(n,:),:);
    V = E(1:3,:);
    [gauss_points_loc, gauss_weights_loc] = gauss_lengdre_loc(V, "2d");
    phi = basis_function_loc(V, gauss_points_loc(:,1), gauss_points_loc(:,2), mesh.basis_type, "x");
    ff1 = pde.f1(gauss_points_loc(:,1), gauss_points_loc(:,2), t);
    ff2 = pde.f2(gauss_points_loc(:,1), gauss_points_loc(:,2), t);
    for i = 1:mesh.Nlb
        f1 = ff1 .* phi(:,i);
        f2 = ff2 .* phi(:,i);
        r1 = gauss_weights_loc * f1;
        r2 = gauss_weights_loc * f2;
        vector_b1(mesh.elem(n,i),1) = vector_b1(mesh.elem(n,i),1) + r1;
        vector_b2(mesh.elem(n,i),1) = vector_b2(mesh.elem(n,i),1) + r2;
    end
end
end